"""
API限流工具
提供基于用户、IP和接口的限流功能
"""
from flask import request, jsonify, g
from functools import wraps
from datetime import datetime, timedelta
from flask import current_app

# Redis是可选的，如果没有安装则使用内存存储
try:
    import redis
    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False

# Redis连接池（如果未配置Redis，使用内存存储）
_rate_limit_store = {}
_redis_client = None

def get_redis_client():
    """获取Redis客户端（可选）"""
    global _redis_client
    if not REDIS_AVAILABLE:
        return None
    
    if _redis_client is None:
        try:
            redis_host = current_app.config.get('REDIS_HOST', 'localhost')
            redis_port = current_app.config.get('REDIS_PORT', 6379)
            redis_db = current_app.config.get('REDIS_DB', 0)
            _redis_client = redis.Redis(
                host=redis_host,
                port=redis_port,
                db=redis_db,
                decode_responses=True
            )
            _redis_client.ping()
        except Exception:
            _redis_client = None
    return _redis_client

def get_rate_limit_key(identifier, endpoint):
    """生成限流键"""
    return f"rate_limit:{endpoint}:{identifier}"

def check_rate_limit(identifier, endpoint, limit, window):
    """
    检查是否超过限流
    
    Args:
        identifier: 标识符（用户ID或IP）
        endpoint: 接口路径
        limit: 限制次数
        window: 时间窗口（秒）
    
    Returns:
        tuple: (is_allowed, remaining, reset_time)
    """
    redis_client = get_redis_client()
    key = get_rate_limit_key(identifier, endpoint)
    
    # 如果使用内存存储，确保_rate_limit_store已初始化
    if not redis_client:
        global _rate_limit_store
        if '_rate_limit_store' not in globals():
            _rate_limit_store = {}
    
    if redis_client:
        # 使用Redis
        current = redis_client.get(key)
        if current is None:
            redis_client.setex(key, window, 1)
            return True, limit - 1, datetime.utcnow() + timedelta(seconds=window)
        
        current_count = int(current)
        if current_count >= limit:
            ttl = redis_client.ttl(key)
            return False, 0, datetime.utcnow() + timedelta(seconds=ttl)
        
        redis_client.incr(key)
        return True, limit - current_count - 1, datetime.utcnow() + timedelta(seconds=redis_client.ttl(key))
    else:
        # 使用内存存储
        now = datetime.utcnow()
        if key not in _rate_limit_store:
            _rate_limit_store[key] = {
                'count': 1,
                'reset_time': now + timedelta(seconds=window)
            }
            return True, limit - 1, _rate_limit_store[key]['reset_time']
        
        record = _rate_limit_store[key]
        if now > record['reset_time']:
            record['count'] = 1
            record['reset_time'] = now + timedelta(seconds=window)
            return True, limit - 1, record['reset_time']
        
        if record['count'] >= limit:
            return False, 0, record['reset_time']
        
        record['count'] += 1
        return True, limit - record['count'], record['reset_time']

def rate_limit(limit=60, window=60, per='ip'):
    """
    装饰器：API限流
    
    Args:
        limit: 限制次数
        window: 时间窗口（秒）
        per: 限流维度 ('ip', 'user', 'endpoint')
    
    Usage:
        @rate_limit(limit=10, window=60, per='user')
        def api_endpoint():
            return jsonify({'message': 'Success'})
    """
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            # 确定标识符
            if per == 'user':
                try:
                    from src.utils.jwt_utils import get_current_user
                    user = get_current_user()
                    identifier = str(user['user_id']) if user else get_client_ip()
                except:
                    identifier = get_client_ip()
            elif per == 'endpoint':
                identifier = request.endpoint
            else:  # per == 'ip'
                identifier = request.remote_addr
            
            endpoint = request.endpoint or request.path
            
            # 检查限流
            is_allowed, remaining, reset_time = check_rate_limit(
                identifier, endpoint, limit, window
            )
            
            if not is_allowed:
                return jsonify({
                    'status': 'error',
                    'message': '请求过于频繁，请稍后再试',
                    'retry_after': int((reset_time - datetime.utcnow()).total_seconds())
                }), 429
            
            # 添加限流信息到响应头
            response = f(*args, **kwargs)
            if hasattr(response, 'headers'):
                response.headers['X-RateLimit-Limit'] = str(limit)
                response.headers['X-RateLimit-Remaining'] = str(remaining)
                response.headers['X-RateLimit-Reset'] = str(int(reset_time.timestamp()))
            
            return response
        return decorated_function
    return decorator

def get_client_ip():
    """获取客户端IP地址"""
    if request.headers.get('X-Forwarded-For'):
        return request.headers.get('X-Forwarded-For').split(',')[0].strip()
    elif request.headers.get('X-Real-IP'):
        return request.headers.get('X-Real-IP')
    else:
        return request.remote_addr

